import re
from typing import Optional

class Alphabet(object):
    ALPHABET= {
        "DNA":"ATGC",
        "RNA":"AUGC",
        "protein":"ACDEFGHIKLMNPQRSTVWY"
    }
    ERROR_ALPHABET_PATTERN = {
        "DNA":re.compile(r'[^ATGCN]'),
        "RNA":re.compile(r'[^AUGCN]'),
        "protein":re.compile(r'[^ACDEFGHIKLMNPQRSTVWY]')
    }
    GEN_CODE = {
        'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
        'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
        'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
        'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
        'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
        'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
        'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
        'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
        'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
        'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
        'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
        'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
        'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
        'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
        'TAC':'Y', 'TAT':'Y', 'TAA':'_', 'TAG':'_',
        'TGC':'C', 'TGT':'C', 'TGA':'_', 'TGG':'W'}

class MySeq(Alphabet):
    def __init__(self,seq:str,seq_type:str="DNA") -> None:
        self.seq = seq.upper()
        self.seq_type = seq_type

    def __len__(self) -> int:
        return len(self.seq)
    
    def __getitem__(self,index) -> str:
        return self.seq[index]
    
    def __getslice__(self,i,j) -> str:
        return self.seq[i:j]
    
    def __str__(self) -> str:
        return self.seq
    
    def get_seq_biotype(self) -> str:
        return self.seq_type
    
    def show_info_seq(self) -> None:
        print(f"Sequence: {self.seq} biotype: {self.seq_type}")
    
    def Alphabet(self) -> str|None:
        return MySeq.ALPHABET.get(self.seq_type,None)
    
    def validate(self) -> bool:
        '''
        raise error when can't pass validation
        '''
        match_iter = re.finditer(
            MySeq.ERROR_ALPHABET_PATTERN[self.seq_type],self.seq)
        
        for error in match_iter:
            raise ValueError(
                "{} at {} index is unsupported in alphabet".format(
                    error.group()[0],error.span()[0]))
        return True
    
    def transcription(self) -> "MySeq":
        if self.seq_type == "DNA":
            return MySeq(self.seq.replace("T","U"),"RNA")
        else:
            raise ValueError(f"the sequence type is {self.seq_type}, expect to DNA")
        
    def reverse_comp(self) -> "MySeq":
        if self.seq_type == "DNA" :
            comp_table = str.maketrans('ACGTacgtRYMKrymkVBHDvbhd', 'TGCAtgcaYRKMyrkmBVDHbvdh')
            seq = self.seq.translate(comp_table)
            return MySeq(seq[::-1],"DNA")
        else:
            raise ValueError(f"the sequence type is {self.seq_type}, expect to DNA")
    
    def translate(self, ini_pos:Optional[int]=0) -> "MySeq":
        if self.seq_type == "DNA" :
            seq = ''.join((self.__translate_codon(self.seq[pos:pos+3]) \
                            for pos in range(ini_pos,len(self.seq)-2,3)))
            return MySeq(seq,"protein")
        elif self.seq_type == "RNA" :
            dna_seq = self.seq.replace('U','T')
            seq = ''.join((self.__translate_codon(dna_seq[pos:pos+3]) \
                            for pos in range(ini_pos,len(dna_seq)-2,3)))
            return MySeq(seq,"protein")
        else:
            raise ValueError(f"the sequence type is {self.seq_type}, expect to DNA")
        
    def __translate_codon(self,codon) -> str:
        return MySeq.GEN_CODE.get(codon,'X')
    

if __name__ == "__main__":
    s1 = MySeq("ATGTGATAAGAATAGAATGCTGAATAAATAGAATGACAT")
    s2 = MySeq("MKVVLSVQERSVVSLL","protein" )
    print(s1.validate(), s2.validate())
    print(s1)
    s3 = s1.transcription()
    s3.show_info_seq()
    s4 = s1.reverse_comp().translate()
    s4.show_info_seq()